#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
从 ROS bag 文件中提取图像并保存为图片文件（支持抽帧）
"""

import os
import sys
import argparse
import cv2
import rosbag
from sensor_msgs.msg import Image
from cv_bridge import CvBridge

def parse_args():
    parser = argparse.ArgumentParser(
        description="从 ROS bag 文件中提取图像并保存为图片文件（支持抽帧）",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("bag_file", help="ROS bag 文件路径")
    parser.add_argument("output_dir", help="保存图片的目录")
    parser.add_argument("image_topic", help="图像话题名称")
    parser.add_argument("--skip", type=int, default=1,
                        help="抽帧间隔：每 N 帧保留 1 帧，必须 ≥1")
    args = parser.parse_args()

    # 参数合法性检查
    if args.skip < 1:
        parser.error("--skip 必须 ≥ 1")
    if not os.path.isfile(args.bag_file):
        parser.error("bag 文件不存在：{}".format(args.bag_file))
    return args

def main():
    args = parse_args()

    # 自动创建输出目录
    os.makedirs(args.output_dir, exist_ok=True)

    bridge = CvBridge()
    count = 0   # 已读帧
    saved = 0   # 已保存帧

    try:
        with rosbag.Bag(args.bag_file, 'r') as bag:
            for topic, msg, t in bag.read_messages(topics=[args.image_topic]):
                count += 1
                if (count - 1) % args.skip != 0:
                    continue

                try:
                    cv_image = bridge.imgmsg_to_cv2(msg, "bgr8")
                except Exception as e:
                    print("警告：转换图像失败 -> {}".format(e), file=sys.stderr)
                    continue

                timestamp = msg.header.stamp
                filename = os.path.join(args.output_dir,
                                        "{}.png".format(timestamp))
                cv2.imwrite(filename, cv_image)
                saved += 1
                print("已保存：{}  ({} / {})".format(filename, saved, count))

    except rosbag.ROSBagException as e:
        print("ROS bag 读取错误：{}".format(e), file=sys.stderr)
        sys.exit(1)

    print("共读取 {} 帧，保存 {} 帧".format(count, saved))

if __name__ == "__main__":
    main()
